
import numpy as np
import torch
from multiprocessing import Pool
import scipy

import joblib

def CalSigmaRho(DisG, DisE, enlarge, v, perplexity, dataset, ndata):
    name=str([enlarge, v, perplexity, dataset, ndata])
    try:
        PE,PG,rhoE,rhoG = joblib.load("save/PG,PE,rhoG,rhoE{}.pkl".format(name))
        print('-----> load PG,PE,rhoG,rhoE from ./save')

    except:
        gamma = CalGammaF(v)
        copyDisG = np.copy(DisG)
        infcopyDisG = np.isposinf(copyDisG)
        copyDisG[infcopyDisG] = -100
        maxDisG = copyDisG.max(axis=0)

        for i in range(DisG.shape[0]):
            DisG[i][infcopyDisG[i]] = maxDisG[i]*enlarge

        print('sigma search {} DisG')

        rhoG, sigmaG = InitSigmaSearchEverySample(
            perplexity=perplexity,
            gamma=gamma,
            v=v,
            dist=DisG,
        )

        PG = CalPt(
            dist=DisG,
            rho=rhoG,
            sigma_array=sigmaG,
            gamma=gamma,
            v=v,
        )

        rhoE, sigmaE = InitSigmaSearchEverySample(
            perplexity=perplexity,
            gamma=gamma,
            v=v,
            dist=DisE,
        )

        PE = CalPt(
            dist=DisE,
            rho=rhoE,
            sigma_array=sigmaE,
            gamma=gamma,
            v=v,
        )
        
        joblib.dump([PE,PG,rhoE,rhoG], "save/PG,PE,rhoG,rhoE{}.pkl".format(name))
        print('-----> save PG,PE,rhoG,rhoE to ./save')
    
    return PG, PE, rhoG, rhoE

def CalGammaF(v : float):
    # for v in vList:
    a = scipy.special.gamma((v + 1) / 2)
    b = np.sqrt(v * np.pi) * scipy.special.gamma(v / 2)
    out=a / b
    return out

def InitSigmaSearchEverySample(perplexity, gamma, v, dist, num_layer=2):

    
    distC = np.copy(dist)
    distC[distC < 1e-11] = 1e16
    rho = np.min(distC, axis=1)

    print("start pool search")

    r = PoolRunner(
        n=dist.shape[0],
        N_NEIGHBOR=perplexity,
        dist=dist,
        rho=rho,
        gamma=gamma,
        v=v,
        pow=2,
    )
    sigmaList = np.array(r.Getout())

    return rho, sigmaList

def CalPt(dist, rho, sigma_array, gamma, v=100, split=1):

    dist_rho = (dist - rho.reshape(-1, 1)) / sigma_array.reshape(-1, 1)
    dist_rho[dist_rho < 0] = 0
    
    Pij = np.power(
        gamma
        * np.power((1 + dist_rho / v), -1 * (v + 1) / 2)
        * np.sqrt(2 * 3.14),
        2,
    )
    P = Pij + Pij.T - np.multiply(Pij, Pij.T)

    return P



class PoolRunner(object):
    def __init__(self, n, N_NEIGHBOR, dist, rho, gamma, v, pow=2):
        pool = Pool(processes=3)

        print(n)
        result = []
        for dist_row in range(n):

            result.append(
                pool.apply_async(
                    sigma_binary_search,
                    (N_NEIGHBOR, dist[dist_row], rho[dist_row], gamma, v, pow),
                )
            )
        print("start calculate sigma")
        pool.close()
        pool.join()
        sigma_array = []
        for i in result:
            sigma_array.append(i.get())
        self.sigma_array = np.array(sigma_array)
        print("\nMean sigma = " + str(np.mean(sigma_array)))
        print("finish calculate sigma")

    def Getout(self,):
        return self.sigma_array


def sigma_binary_search(fixed_k, dist_row_line, rho_line, gamma, v, pow=2):
    """
    Solve equation k_of_sigma(sigma) = fixed_k
    with respect to sigma by the binary search algorithm
    """
    sigma_lower_limit = 0
    sigma_upper_limit = 100
    for i in range(20):
        approx_sigma = (sigma_lower_limit + sigma_upper_limit) / 2
        k_value = func(approx_sigma, dist_row_line, rho_line, gamma, v, pow=pow)
        if k_value < fixed_k:
            sigma_lower_limit = approx_sigma
        else:
            sigma_upper_limit = approx_sigma
        if np.abs(fixed_k - k_value) <= 1e-4:
            break
    return approx_sigma


def func(sigma, dist_row_line, rho_line, gamma, v, pow=2):
    d = (dist_row_line - rho_line) / sigma
    d[d < 0] = 0
    p = np.power(
        gamma * np.power((1 + d / v), -1 * (v + 1) / 2) * np.sqrt(2 * 3.14), pow
    )
    return np.power(2, np.sum(p))